import numpy as np
from PIL import Image, ImageDraw, ImageFont
import os
import glob
import math


class AdaptiveScaleBar:
    def __init__(self, resolution_px=300, resolution_um=20):
        """初始化比例尺系统

        Args:
            resolution_px: 图像中的像素数
            resolution_um: 对应的实际长度(微米)
        """
        self.px_to_um = resolution_um / resolution_px  # 每像素对应的微米数
        self.standard_scales = [1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000, 5000]
        self.unit_thresholds = {1000: ("mm", 1000)}

        # 存储分辨率信息，用于自适应调整线宽
        self.resolution_px = resolution_px
        self.resolution_um = resolution_um

    def _select_best_scale(self, img_width, min_bar_px=30):
        """选择最佳比例尺长度，确保最小像素宽度"""
        ideal_bar_width_px = img_width / 6
        ideal_bar_width_px = max(ideal_bar_width_px, min_bar_px)
        ideal_bar_width_um = ideal_bar_width_px * self.px_to_um
        best_scale = min(self.standard_scales, key=lambda x: abs(x - ideal_bar_width_um))
        return best_scale

    def _format_scale_text(self, scale_um):
        """格式化比例尺文本，自动转换单位"""
        for threshold, (unit, divisor) in self.unit_thresholds.items():
            if scale_um >= threshold:
                return f"{scale_um / divisor:.0f} {unit}"
        return f"{scale_um:.0f} μm"

    def _parse_color(self, color):
        """解析颜色参数"""
        if isinstance(color, str):
            # 预定义颜色
            color_map = {
                'red': (255, 0, 0, 255),
                'green': (0, 255, 0, 255),
                'blue': (0, 0, 255, 255),
                'white': (255, 255, 255, 255),
                'black': (0, 0, 0, 255),
                'yellow': (255, 255, 0, 255),
                'cyan': (0, 255, 255, 255),
                'magenta': (255, 0, 255, 255),
                'orange': (255, 165, 0, 255)
            }

            if color.lower() in color_map:
                return color_map[color.lower()]

            # 十六进制颜色
            if color.startswith('#'):
                color = color.lstrip('#')
                if len(color) == 6:
                    r = int(color[0:2], 16)
                    g = int(color[2:4], 16)
                    b = int(color[4:6], 16)
                    return (r, g, b, 255)
                elif len(color) == 8:
                    r = int(color[0:2], 16)
                    g = int(color[2:4], 16)
                    b = int(color[4:6], 16)
                    a = int(color[6:8], 16)
                    return (r, g, b, a)

        # 如果是元组，直接返回
        if isinstance(color, tuple) and len(color) >= 3:
            if len(color) == 3:
                return color + (255,)  # 添加透明度
            return color

        # 默认返回红色
        return (255, 0, 0, 255)

    def _calculate_adaptive_line_width(self, img_width, scale_um):
        """计算自适应线宽

        基于图像尺寸和物理分辨率计算合适的线宽，确保比例尺在不同分辨率下都清晰可见
        """
        # 计算比例尺在图像中所占比例
        scale_px = scale_um / self.px_to_um
        scale_ratio = scale_px / img_width

        # 根据比例尺在图像中的比例计算线宽
        # 当比例尺占图像宽度比例较小时，增加线宽
        if scale_ratio < 0.1:
            base_width = max(2, int(img_width * 0.003))
        else:
            base_width = max(2, int(img_width * 0.002))

        # 根据物理分辨率调整线宽
        # 当物理分辨率较小(如20μm)时，增加线宽以确保可见性
        resolution_factor = max(1.0, 50.0 / self.resolution_um)

        # 最终线宽，确保至少为2像素
        line_width = max(2, int(base_width * resolution_factor))

        return line_width

    def add_scale_bar(self, image_path, output_path=None,
                      line_width=None, font_size=None,
                      color=(255, 0, 0, 255),  # 红色
                      position="bottom-left"):
        """向图像添加红色线段比例尺

        Args:
            image_path: 输入图像路径
            output_path: 输出图像路径
            line_width: 线段宽度(像素)，None为自动计算
            font_size: 字体大小(像素)，None为自动
            color: 比例尺颜色，可以是(R,G,B,A)元组、颜色名称或十六进制字符串
            position: 位置，"bottom-left", "bottom-right", "top-left", "top-right"
        """
        # 解析颜色
        color = self._parse_color(color)

        # 打开图像并确保为RGBA模式(支持透明度)
        img = Image.open(image_path)
        if img.mode != 'RGBA':
            img = img.convert('RGBA')

        img_width, img_height = img.size

        # 自动计算参数
        is_small_image = img_width < 400 or img_height < 400
        margin = max(10, int(img_width * 0.02))  # 至少10像素边距

        # 选择最佳比例尺
        min_bar_px = max(30, int(img_width * 0.1))  # 确保比例尺足够可见
        scale_um = self._select_best_scale(img_width, min_bar_px)
        scale_px = int(round(scale_um / self.px_to_um))  # 取整数像素
        scale_text = self._format_scale_text(scale_um)

        # 自动计算线宽
        if line_width is None:
            line_width = self._calculate_adaptive_line_width(img_width, scale_um)

        # 自动计算字体大小
        if font_size is None:
            # 字体大小与线宽相关，确保比例协调
            font_size = max(10, int(line_width * 3.5))
            if is_small_image:
                font_size = max(8, font_size)  # 小图像最小8pt

        # 创建透明叠加层
        overlay = Image.new('RGBA', img.size, (0, 0, 0, 0))
        draw = ImageDraw.Draw(overlay)

        # 设置字体
        try:
            # 尝试使用无衬线字体，更清晰
            font_paths = [
                "arial.ttf",
                "DejaVuSans.ttf",
                "NotoSans-Regular.ttf",
                "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"
            ]

            font = None
            for path in font_paths:
                try:
                    font = ImageFont.truetype(path, font_size)
                    break
                except IOError:
                    continue

            if font is None:
                font = ImageFont.load_default()

        except Exception:
            font = ImageFont.load_default()

        # 计算文本尺寸
        try:
            # PIL 9.0.0+
            bbox = draw.textbbox((0, 0), scale_text, font=font)
            text_width = bbox[2] - bbox[0]
            text_height = bbox[3] - bbox[1]
        except AttributeError:
            # 旧版PIL
            text_width = draw.textlength(scale_text, font=font)
            text_height = font_size

        # 根据位置确定起始坐标
        if position == "bottom-left":
            start_x = margin
            start_y = img_height - margin
        elif position == "bottom-right":
            start_x = img_width - margin - scale_px
            start_y = img_height - margin
        elif position == "top-left":
            start_x = margin
            start_y = margin
        elif position == "top-right":
            start_x = img_width - margin - scale_px
            start_y = margin
        else:
            raise ValueError("位置参数无效，请使用: bottom-left, bottom-right, top-left, top-right")

        end_x = start_x + scale_px
        end_y = start_y

        # 绘制线段比例尺 - 红色线段
        draw.line(
            [(start_x, start_y), (end_x, end_y)],
            fill=color,
            width=line_width
        )

        # 添加短垂线
        tick_height = max(4, line_width * 2)
        draw.line([(start_x, start_y - tick_height // 2), (start_x, start_y + tick_height // 2)], fill=color,
                  width=line_width)
        draw.line([(end_x, start_y - tick_height // 2), (end_x, start_y + tick_height // 2)], fill=color,
                  width=line_width)

        # 计算文本位置
        text_x = start_x + (scale_px - text_width) / 2

        # 根据位置调整文本Y坐标
        if position.startswith("bottom"):
            text_y = start_y - text_height - 5
        else:
            text_y = end_y + 5

        # 绘制文本
        draw.text(
            (text_x, text_y),
            scale_text,
            font=font,
            fill=color
        )

        # 合并图层
        result = Image.alpha_composite(img, overlay)

        # 保存结果
        if output_path is None:
            base, ext = os.path.splitext(image_path)
            output_path = f"{base}_scale{ext}"

        # 确保输出格式正确
        if output_path.lower().endswith('.jpg'):
            result = result.convert('RGB')

        result.save(output_path)
        return output_path

    def process_batch(self, input_dir, pattern="*.jpg", output_dir=None,
                      recursive=False, overwrite=False, **kwargs):
        """批量处理图像添加比例尺

        Args:
            input_dir: 输入图像目录
            pattern: 文件匹配模式，如'*.jpg'或['*.jpg', '*.png']
            output_dir: 输出图像目录，None表示覆盖原图像
            recursive: 是否递归处理子目录
            overwrite: 是否覆盖原文件
            **kwargs: 传递给add_scale_bar的参数

        Returns:
            处理的文件列表
        """
        # 确保输出目录存在
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)

        processed_files = []

        # 获取所有匹配的文件
        all_files = []
        if isinstance(pattern, str):
            pattern = [pattern]

        for pat in pattern:
            if recursive:
                # 递归搜索所有子目录
                for root, _, _ in os.walk(input_dir):
                    files = glob.glob(os.path.join(root, pat))
                    all_files.extend(files)
            else:
                # 只搜索当前目录
                files = glob.glob(os.path.join(input_dir, pat))
                all_files.extend(files)

        # 去重
        all_files = sorted(list(set(all_files)))

        # 处理每个文件
        for i, file_path in enumerate(all_files):
            try:
                # 显示进度
                print(f"处理 {i + 1}/{len(all_files)}: {os.path.basename(file_path)}")

                # 确定输出路径
                if output_dir:
                    # 保持相对路径结构
                    rel_path = os.path.relpath(file_path, input_dir)
                    out_path = os.path.join(output_dir, rel_path)
                    # 确保输出子目录存在
                    os.makedirs(os.path.dirname(out_path), exist_ok=True)

                    # 如果不覆盖，则添加后缀
                    if not overwrite:
                        base, ext = os.path.splitext(out_path)
                        out_path = f"{base}_scale{ext}"
                else:
                    # 原地覆盖或添加后缀
                    if overwrite:
                        out_path = file_path
                    else:
                        base, ext = os.path.splitext(file_path)
                        out_path = f"{base}_scale{ext}"

                # 处理图像
                result_path = self.add_scale_bar(file_path, out_path, **kwargs)
                processed_files.append(result_path)

            except Exception as e:
                print(f"处理文件 {file_path} 时出错: {e}")

        return processed_files

    def upscale_and_add_scale(self, image_path, output_path=None, scale_factor=2, **kwargs):
        """先放大图像再添加比例尺，然后缩回原尺寸"""
        # 打开图像
        img = Image.open(image_path)
        original_size = img.size

        # 放大图像
        new_size = (original_size[0] * scale_factor, original_size[1] * scale_factor)
        img_upscaled = img.resize(new_size, Image.LANCZOS)

        # 保存临时放大图像
        temp_path = f"{os.path.splitext(image_path)[0]}_temp.png"
        img_upscaled.save(temp_path)

        # 计算自适应线宽，如果未指定
        if 'line_width' not in kwargs:
            # 选择最佳比例尺
            min_bar_px = max(30, int(new_size[0] * 0.1))
            scale_um = self._select_best_scale(new_size[0], min_bar_px)
            # 计算线宽
            line_width = self._calculate_adaptive_line_width(new_size[0], scale_um)
            kwargs['line_width'] = line_width

        # 添加比例尺到放大图像
        self.add_scale_bar(
            temp_path,
            temp_path,
            font_size=kwargs.get('font_size', None),
            color=kwargs.get('color', (255, 0, 0, 255)),
            position=kwargs.get('position', 'bottom-left'),
            **{k: v for k, v in kwargs.items() if k not in ['font_size', 'color', 'position']}
        )

        # 读取添加了比例尺的图像并缩小回原尺寸
        img_with_scale = Image.open(temp_path)
        img_final = img_with_scale.resize(original_size, Image.LANCZOS)

        # 保存结果
        if output_path is None:
            base, ext = os.path.splitext(image_path)
            output_path = f"{base}_scale{ext}"

        img_final.save(output_path)

        # 删除临时文件
        try:
            os.remove(temp_path)
        except:
            pass

        return output_path

    def process_batch_upscale(self, input_dir, pattern="*.jpg", output_dir=None,
                              recursive=False, overwrite=False, scale_factor=2, **kwargs):
        """批量处理图像，使用超分辨率技术添加比例尺"""
        # 确保输出目录存在
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)

        processed_files = []

        # 获取所有匹配的文件
        all_files = []
        if isinstance(pattern, str):
            pattern = [pattern]

        for pat in pattern:
            if recursive:
                # 递归搜索所有子目录
                for root, _, _ in os.walk(input_dir):
                    files = glob.glob(os.path.join(root, pat))
                    all_files.extend(files)
            else:
                # 只搜索当前目录
                files = glob.glob(os.path.join(input_dir, pat))
                all_files.extend(files)

        # 去重
        all_files = sorted(list(set(all_files)))

        # 处理每个文件
        for i, file_path in enumerate(all_files):
            try:
                # 显示进度
                print(f"处理 {i + 1}/{len(all_files)}: {os.path.basename(file_path)} (超分辨率)")

                # 确定输出路径
                if output_dir:
                    # 保持相对路径结构
                    rel_path = os.path.relpath(file_path, input_dir)
                    out_path = os.path.join(output_dir, rel_path)
                    # 确保输出子目录存在
                    os.makedirs(os.path.dirname(out_path), exist_ok=True)

                    # 如果不覆盖，则添加后缀
                    if not overwrite:
                        base, ext = os.path.splitext(out_path)
                        out_path = f"{base}_scale{ext}"
                else:
                    # 原地覆盖或添加后缀
                    if overwrite:
                        out_path = file_path
                    else:
                        base, ext = os.path.splitext(file_path)
                        out_path = f"{base}_scale{ext}"

                # 处理图像
                result_path = self.upscale_and_add_scale(
                    file_path,
                    out_path,
                    scale_factor=scale_factor,
                    **kwargs
                )
                processed_files.append(result_path)

            except Exception as e:
                print(f"处理文件 {file_path} 时出错: {e}")

        return processed_files


# 使用示例
if __name__ == "__main__":
    # 创建比例尺系统
    scale_system = AdaptiveScaleBar(resolution_px=300, resolution_um=20)

    # 单图处理示例
    # print("==== 单图处理 ====")
    # result_path = scale_system.add_scale_bar(
    #     "1_03_1.jpg",  # 输入图像
    #     color="red",  # 红色，可以是颜色名称、十六进制(#FF0000)或RGB元组
    #     position="bottom-left"  # 位置
    #     # 不指定线宽，使用自适应计算
    # )
    # print(f"比例尺已添加，保存至: {result_path}")

    # 批量处理示例
    print("\n==== 批量处理 ====")
    # 获取当前目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    input_dir = current_dir  # 当前目录
    output_dir = os.path.join(current_dir, "output")  # 输出目录

    # 创建输出目录(如果不存在)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 批量处理所有JPG图像
    processed_files = scale_system.process_batch(
        input_dir=input_dir,
        pattern=["*.jpg", "*.jpeg"],  # 处理所有JPG图像
        output_dir=output_dir,  # 输出到output子目录
        color="red",  # 红色
        position="bottom-left",  # 位置
        recursive=False  # 不递归处理子目录
        # 不指定线宽，使用自适应计算
    )
    print(f"批量处理完成，共处理 {len(processed_files)} 个文件。")

    # 批量处理小尺寸图像(使用超分辨率技术)
    print("\n==== 批量处理小尺寸图像 ====")
    small_img_dir = os.path.join(current_dir, "small_images")  # 假设有small_images子目录
    # 如果目录存在，则处理
    if os.path.exists(small_img_dir):
        processed_small = scale_system.process_batch_upscale(
            input_dir=small_img_dir,
            pattern="*.jpg",
            output_dir=os.path.join(output_dir, "small"),
            scale_factor=3,  # 放大3倍
            color="white"  # 白色
            # 不指定线宽，使用自适应计算
        )
        print(f"小尺寸图像处理完成，共处理 {len(processed_small)} 个文件。")
    else:
        print(f"小尺寸图像目录 {small_img_dir} 不存在，跳过处理。")
